#include <algorithm>
#include <cmath>
#include <iostream>
#include <limits>
#include <vector>


using namespace std;

// 辅助函数：计算窗口内均值和方差（带边界处理）
template <typename T>
void computeWindowStats(
    const vector<vector<T>>& img,
    int x,
    int y,
    int radius,
    double& mean,
    double& variance,
    const vector<vector<double>>& meanI,
    const vector<vector<double>>& covI
)
{
    int channels = img[0].size();
    mean = 0.0;
    variance = 0.0;

    int half = radius * 2 + 1;
    double sumMean = 0.0, sumVar = 0.0;

    // 遍历窗口区域
    for (int i = -radius; i <= radius; ++i) {
        for (int j = -radius; j <= radius; ++j) {
            int xi = x + i;
            int yj = y + j;

            // 边界处理：镜像填充
            if (xi < 0) {
                xi = -xi;
            }
            if (xi >= img.size()) {
                xi = 2 * (img.size() - 1) - xi;
            }
            if (yj < 0) {
                yj = -yj;
            }
            if (yj >= img[0].size()) {
                yj = 2 * (img[0].size() - 1) - yj;
            }

            double val = img[xi][yj];
            sumMean += val;
        }
    }
    mean = sumMean / (half * half);

    // 计算方差（假设meanI和covI已预计算）
    sumVar = covI[x][y] - (meanI[x][y] * meanI[x][y]);
}

vector<vector<double>> guidedFilter(
    const vector<vector<double>>& input, // 输入图像（可多通道）
    const vector<vector<double>>& guide, // 引导图像（通常为原图）
    int radius,                          // 滤波半径
    double epsilon                       // 正则化参数
)
{
    if (input.empty() || guide.empty()) {
        return {};
    }
    if (input.size() != guide.size() || (input[0].size() != guide[0].size())) {
        throw invalid_argument("输入图像和引导图像尺寸必须一致");
    }

    int rows = input.size();
    int cols = input[0].size();
    int channels = input[0].size();

    // 预计算引导图像的均值和协方差
    vector<vector<double>> meanI(rows, vector<double>(cols, 0.0));
    vector<vector<double>> meanP(rows, vector<double>(cols, 0.0));
    vector<vector<double>> corrI(rows, vector<double>(cols, 0.0));
    vector<vector<double>> corrIP(rows, vector<double>(cols, 0.0));

    // 使用积分图像加速统计量计算（此处简化为直接遍历）
    for (int i = 0; i < rows; ++i) {
        for (int j = 0; j < cols; ++j) {
            double sumI = 0.0, sumP = 0.0, sumII = 0.0, sumIP = 0.0;
            int half = radius * 2 + 1;

            // 遍历窗口
            for (int x = -radius; x <= radius; ++x) {
                for (int y = -radius; y <= radius; ++y) {
                    int xi = i + x;
                    int yj = j + y;

                    // 边界处理：镜像填充
                    if (xi < 0) {
                        xi = -xi;
                    }
                    if (xi >= rows) {
                        xi = 2 * (rows - 1) - xi;
                    }
                    if (yj < 0) {
                        yj = -yj;
                    }
                    if (yj >= cols) {
                        yj = 2 * (cols - 1) - yj;
                    }

                    double I = guide[xi][yj];
                    double P = input[xi][yj];

                    sumI += I;
                    sumP += P;
                    sumII += I * I;
                    sumIP += I * P;
                }
            }

            meanI[i][j] = sumI / (half * half);
            meanP[i][j] = sumP / (half * half);
            corrI[i][j] = sumII / (half * half) - meanI[i][j] * meanI[i][j];
            corrIP[i][j] = sumIP / (half * half) - meanI[i][j] * meanP[i][j];
        }
    }

    // 计算线性系数 a 和 b
    vector<vector<double>> a(rows, vector<double>(cols, 0.0));
    vector<vector<double>> b(rows, vector<double>(cols, 0.0));

    for (int i = 0; i < rows; ++i) {
        for (int j = 0; j < cols; ++j) {
            double varI = corrI[i][j] + epsilon;
            double covIP = corrIP[i][j];
            a[i][j] = covIP / varI;
            b[i][j] = meanP[i][j] - a[i][j] * meanI[i][j];
        }
    }

    // 对 a 和 b 进行均值滤波
    auto boxFilter = [](const vector<vector<double>>& src, int radius) {
        int half = radius * 2 + 1;
        int rows = src.size(), cols = src[0].size();
        vector<vector<double>> dst(rows, vector<double>(cols, 0.0));
        vector<vector<double>> sum(rows + 1, vector<double>(cols + 1, 0.0));

        // 积分图像加速
        for (int i = 0; i < rows; ++i) {
            for (int j = 0; j < cols; ++j) {
                sum[i + 1][j + 1] = src[i][j] + sum[i][j + 1] + sum[i + 1][j] - sum[i][j];
            }
        }

        for (int i = 0; i < rows; ++i) {
            for (int j = 0; j < cols; ++j) {
                int x1 = max(0, i - radius);
                int y1 = max(0, j - radius);
                int x2 = min(rows - 1, i + radius);
                int y2 = min(cols - 1, j + radius);
                dst[i][j] =
                    sum[x2 + 1][y2 + 1] - sum[x1][y2 + 1] - sum[x2 + 1][y1] + sum[x1][y1];
                dst[i][j] /= (x2 - x1 + 1) * (y2 - y1 + 1);
            }
        }

        return dst;
    };

    auto meanA = boxFilter(a, radius);
    auto meanB = boxFilter(b, radius);

    // 计算输出图像
    vector<vector<double>> output(rows, vector<double>(cols, 0.0));
    for (int i = 0; i < rows; ++i) {
        for (int j = 0; j < cols; ++j) {
            output[i][j] = meanA[i][j] * guide[i][j] + meanB[i][j];
        }
    }

    return output;
}

// // 测试用灰度图像类型
// using GrayImage = vector<vector<double>>;

// // 示例：生成测试图像（含噪声）
// GrayImage createTestGrayImage(int rows, int cols)
// {
//     GrayImage img(rows, vector<double>(cols, 128.0));
//     for (int i = 0; i < rows; ++i) {
//         for (int j = 0; j < cols; ++j) {
//             if (rand() % 100 < 10) {            // 10%噪声概率
//                 img[i][j] = 0.0 + rand() % 256; // 添加椒盐噪声
//             }
//         }
//     }
//     return img;
// }

// // 示例：打印图像（调试用）
// void printImage(const GrayImage& img) {
//     for (const auto& row : img) {
//         for (double val : row) {
//             cout << val << " ";
//         }
//         cout << endl;
//     }
// }

// int main() {
//     // 测试灰度图像
//     GrayImage noisyImg = createTestGrayImage(5, 5);
//     cout << "Noisy Image:" << endl;
//     printImage(noisyImg);

//     // 应用导向滤波（使用原图作为引导）
//     auto filteredImg = guidedFilter(noisyImg, noisyImg, 2, 0.01);

//     cout << "\nFiltered Image:" << endl;
//     printImage(filteredImg);

//     return 0;
// }
